#ifndef COARSEN_HEAP_HPP
#define COARSEN_HEAP_HPP
#include "utils.hpp"
#include "coarsen.hpp"
#include <unordered_map>
#include <algorithm>
#include <climits>

struct Elem
{
    int id = -1;// 点编号
    int cnt = 0;// 计数
    Elem() {}
    Elem(int _id, int _cnt): id(_id), cnt(_cnt) {}
    bool operator == (const Elem & b) const {
        return id == b.id;
    }
    bool operator < (const Elem & b) const {
        if (cnt == b.cnt) return id < b.id;
        else return cnt < b.cnt;
    }
};

template<typename idx_t>
class MinHeap
{
private:
    bool assembled = false;
    idx_t capa = 0;
    Elem * heapDataVec = nullptr;
    idx_t * posit_table = nullptr;
public:
    idx_t numCounts = 0;
    MinHeap() {}
    MinHeap(const idx_t num) { initData(num, nullptr); }
    void printPos() {
        for (idx_t i = 0; i < capa; i++) {
            printf(" %3d ", posit_table[i]);
        }
        printf("\n");
    }
    void insert(Elem elem) {
        assert(numCounts < capa);
        assert(elem.id == numCounts);
        heapDataVec [numCounts] = elem;
        posit_table [elem.id] = numCounts;
        numCounts ++;
        percolateUp(numCounts - 1, heapDataVec[numCounts - 1]);// 对该元素做上滤
    }
    void percolateUp(idx_t index, Elem value) {
        idx_t parentIndex = (index - 1) >> 1;
        while (index > 0 && value < heapDataVec[parentIndex]) {// 如果当前位置(由index指示)还不是堆顶，且父节点的值比本位置要大，则继续
            idx_t nameParent = heapDataVec[parentIndex].id;
            posit_table[nameParent] = index;// 更新原父节点的位置
            heapDataVec [index] = heapDataVec [parentIndex];
            index = parentIndex;
            parentIndex = (parentIndex - 1) >> 1;
        }
        // 找到了最终放置的位置
        heapDataVec[index] = value;
        posit_table[value.id] = index;
    }
    Elem delMin() {
        Elem value = heapDataVec[numCounts - 1];// 原末尾元素
        Elem ret   = heapDataVec[0];// 堆顶元素，准备删除
        idx_t name_0 = ret.id; assert(posit_table[name_0] == 0);
        heapDataVec[numCounts - 1] = heapDataVec[0];// 原堆顶元素放到最后（即无效）
        posit_table[name_0] = numCounts - 1;// 记为

        numCounts --;// 堆中元素个数减一
        percolateDown(0, value);// 下滤，value是替换后的新的根节点值
        return ret;
    }
    void percolateDown(idx_t childTree, Elem value) {
        // printf("percoDown id %d cnt %d\n", value.id, value.cnt);
        idx_t index = childTree;// 目标节点索引
        idx_t maxChid = (index + 1) << 1;// 目标节点右子节点
        bool godown = true;// 循环下滤退出标志
        while (godown && maxChid < numCounts) {// 右子节点尚还处于堆内
            godown = false;
            // printf("  maxChid %d     while [0] = id %d cnt %d\n", maxChid, heapDataVec[0].id, heapDataVec[0].cnt);
            if (heapDataVec[maxChid - 1] < heapDataVec[maxChid]) {// 若左孩子小，则更新为左孩子
                maxChid --;
                // printf("  maxChid %d     while [0] = id %d cnt %d\n", maxChid, heapDataVec[0].id, heapDataVec[0].cnt);
            }
            if (heapDataVec[maxChid] < value) {// 若父节点大于孩子，则下滤，此时孩子的位置记录需要更新
                godown = true;
                idx_t nameChid = heapDataVec[maxChid].id;
                posit_table[nameChid] = index;// 更新孩子的位置
                heapDataVec[index] = heapDataVec[maxChid];// 令较小值为交换值（将孩子上移）
                index = maxChid;// 该交换节点索引下移  
                maxChid = (maxChid + 1) << 1;// 重新计算交换节点右子节点  
            }
        }// while loop
        if (maxChid == numCounts) {// 右子节点已经不在堆内了
            if (heapDataVec[maxChid - 1] < value) {// 令左子节点值为交换值
                idx_t nameChid = heapDataVec[maxChid - 1].id;
                posit_table[nameChid] = index;// 更新孩子的位置
                heapDataVec[index] = heapDataVec[maxChid - 1];
                index = maxChid - 1;
            }
        }
        // printf("  final %d     while [0] = id %d cnt %d\n", index, heapDataVec[0].id, heapDataVec[0].cnt);
        // 找到了最终放置的位置
        heapDataVec[index] = value;// 将调整值赋予交换节点
        posit_table[value.id] = index;// 更新元素位置
    }
    void makeHeap() {// 建堆的过程就是一个不断调整堆的过程，循环调用函数percolateDown调整子树  
        if (numCounts < 2) return;
        // 第一个需要调整的子树的根节点地址
        idx_t parent = (numCounts >> 1) - 1;
        while (true) {
            percolateDown(parent, heapDataVec[parent]);
            if (0 == parent) return; // 到达根节点，返回
            -- parent;
        }
    }
    Elem heapTop() {
        return heapDataVec[0]; 
    }
    void initData(const idx_t _num, const Elem * _data) {
        assert(numCounts == 0 && capa == 0);
        numCounts = capa = _num;
        heapDataVec = new Elem [numCounts];
        posit_table = new idx_t[numCounts];
        if (_data == nullptr) {
            for (idx_t j = 0; j < numCounts; j++) {
                heapDataVec[j].id = j;
                posit_table[j] = j;
            }
        } else {
            for (idx_t j = 0; j < numCounts; j++) {
                heapDataVec[j] = _data[j];
                assert(_data[j].id < numCounts);
                posit_table[_data[j].id] = j;// 记录位于数据向量中的哪个位置
            }
        }
    }
    Elem * data() { return heapDataVec; }
    void assemble() {
        assert(assembled == false);
        makeHeap();
        assembled = true;
    }
    void decre(const idx_t num, const idx_t * names) {
        for (idx_t j = 0; j < num; j++) {
            const idx_t name = names[j];
            const idx_t posi = posit_table[name];
            // printf("decre %d at %d\n", name, posi);
            if (posi >= numCounts) continue;// 这个邻居之前已经被处理过了

            assert(heapDataVec[posi].id == name);
            heapDataVec[posi].cnt --;// 计数减一
            percolateUp(posi, heapDataVec[posi]);// 然后做上滤
        }
    }
    // 删除特定序号(名字)的元素
    // 通过将其计数设为INT_MIN(最小的计数值)，然后上滤，则一定会出现在堆顶，再从堆顶删掉
    void delName(const idx_t name) {
        const idx_t posi = posit_table[name]; assert(posi < numCounts); assert(heapDataVec[posi].id == name);
        heapDataVec[posi].cnt = INT_MIN;
        percolateUp(posi, heapDataVec[posi]);
        Elem toDel = delMin(); assert(toDel.id == name);// 确认从堆顶删除的确实是需要删除的那个元素
    }
    ~MinHeap() {
        delete [] heapDataVec; heapDataVec = nullptr;
        delete [] posit_table; posit_table= nullptr;
        numCounts = capa = 0;
    }
};

template<typename idx_t, typename data_t, typename calc_t> static
CoarsenWrapper<idx_t, data_t, calc_t>* pairwise_coarsen(const MyParCSRMatrix<idx_t, data_t, calc_t> & fine_mat,
    const calc_t threshold, const calc_t alpha)
{
    assert(threshold > 0.0);
    int my_pid; MPI_Comm_rank(fine_mat.comm, &my_pid);
    int nprocs; MPI_Comm_size(fine_mat.comm, &nprocs);
    if (my_pid == 0) printf("Exact coarsen\n");
    std::vector<double> record;
    record.emplace_back(MPI_Wtime());

    const idx_t num_fine = fine_mat.end_row - fine_mat.beg_row;// 本进程在细网格所负责的行数
    MinHeap<idx_t> minheap(num_fine);
    Elem * ST_cnts = minheap.data();

    const MyCSRMatrix<idx_t, data_t, calc_t> & fine_diag = *(fine_mat.diag);
    const MyCSRMatrix<idx_t, data_t, calc_t> & fine_offd = *(fine_mat.offd);
    idx_t * loc_map_f2c = new idx_t [num_fine];// 记录本地的第i个细点被映射为粗点的序号

    idx_t thread_flags [omp_get_max_threads()];
    idx_t thread_shared[omp_get_max_threads()];
    for (idx_t t = 0; t < omp_get_max_threads(); t++) thread_flags[t] = -1;
    idx_t * S_up_ptr = new idx_t [num_fine + 1], * S_upstrm = nullptr;
    S_up_ptr[0] = 0;
    #pragma omp parallel
    {
        int tid = omp_get_thread_num();
        int nt  = omp_get_num_threads();
        int tbeg = 0, tend = 0;
        thread_load_balance(num_fine, nt, tid, tbeg, tend);

        idx_t * tl_rpt = new idx_t [tend - tbeg + 1]; tl_rpt[0] = 0;
        std::vector<idx_t> tl_S_upstrm;
        for (idx_t i = tbeg; i < tend; i++) {
            loc_map_f2c[i] = -1;
            const idx_t num_diag = fine_diag.row_ptr[i + 1] - fine_diag.row_ptr[i];
            // 决定某个邻居是否强影响我
            data_t min_neg = 0.0;
            // 检查对角部分的邻居
            for (idx_t p = fine_diag.row_ptr[i]; p < fine_diag.row_ptr[i + 1]; p++) {
                if (fine_diag.col_idx[p] == i) {
                    assert(fine_diag.vals[p] > 0.0);// 对角元要求为正
                    continue;// 对角元跳过
                }
                if (fine_diag.vals[p] < 0.0) {// 原算法好像只对负的非对角元处理？
                    min_neg = MIN(min_neg, fine_diag.vals[p]);
                }
            }
            // 检查非对角部分的邻居
            for (idx_t p = fine_offd.row_ptr[i]; p < fine_offd.row_ptr[i + 1]; p++) {
                if (fine_offd.vals[p] < 0.0) {// 原算法好像只对负的非对角元处理？
                    min_neg = MIN(min_neg, fine_offd.vals[p]);
                }
            }
            // 定义本点受其他点影响强弱的阈值（是一个负数）
            const data_t beta = threshold * min_neg; assert(beta < 0.0);
            data_t vals[num_diag];
            idx_t cnt = 0;
            // 根据阈值，确定强连通的邻居
            // 注意只需要在对角矩阵部分找强连通邻居即可，非对角部分 (non-local) 不用管
            for (idx_t p = fine_diag.row_ptr[i]; p < fine_diag.row_ptr[i + 1]; p++) {
                const idx_t ngb_id = fine_diag.col_idx[p]; assert(ngb_id < num_fine);
                if (ngb_id == i) continue;
                // if (nprocs == 1) if (ngb_id / (num_fine/2) != i / (num_fine/2)) continue;// 模拟多个进程的情况
                if (fine_diag.vals[p] < beta) {
                    tl_S_upstrm.push_back(ngb_id);// S_upstrm.push_back(ngb_id);
                    vals [cnt++] = fine_diag.vals[p];
                    __atomic_fetch_add( & ST_cnts[ngb_id].cnt, 1, __ATOMIC_RELAXED);// ST_cnts[ngb_id].cnt ++;
                }
            }
            const idx_t pend = tl_S_upstrm.size(),
                        pbeg = pend - cnt;
            idx_t * idxs = tl_S_upstrm.data() + pbeg;
            // 根据非零元的(负)值从小到大（即绝对值从大到小）冒泡排序 S_upstrm[pbeg, pend) ~ vals[0, cnt)
            bubble_sort<data_t, idx_t>(cnt, vals, idxs);

            // 记录这个点有多少个强影响的邻居
            tl_rpt[i - tbeg + 1] = pend;// S_up_ptr.push_back(S_upstrm.size());
        }
        assert(tl_rpt[tend - tbeg] == (idx_t) tl_S_upstrm.size());
        
        // 多线程数据汇总
        while (tid > 0 && __atomic_load_n( & thread_flags[tid - 1], __ATOMIC_ACQUIRE) < 1) {  }// 等待上一个线程完成
        idx_t prev_tsum = (tid > 0) ? thread_shared[tid - 1] : 0;
        thread_shared[tid] = prev_tsum + tl_rpt[tend - tbeg];// 本线程内的求和存入多线程共享区
        __atomic_store_n( & thread_flags[tid], 1, __ATOMIC_RELEASE);

        for (idx_t i = tbeg; i < tend; i++) {
            S_up_ptr[i + 1] = prev_tsum + tl_rpt[i - tbeg + 1];
        }
        __atomic_store_n( & thread_flags[tid], 2, __ATOMIC_RELEASE);
        
        if (tid == 0) {
            while (__atomic_load_n( & thread_flags[nt - 1], __ATOMIC_ACQUIRE) < 2) {  }// 等待最后一个线程填充完
            const idx_t nnz = S_up_ptr[num_fine];// printf("nnz %d\n", nnz);
            S_upstrm = new idx_t [nnz];
            __atomic_store_n( & thread_flags[tid], 3, __ATOMIC_RELEASE);
        }
        
        while (tid > 0 && __atomic_load_n( & thread_flags[tid - 1], __ATOMIC_ACQUIRE) < 2) {  }// 等待上一个线程把假定的(assum)行偏移填充完
        while (tid > 0 && __atomic_load_n( & thread_flags[  0    ], __ATOMIC_ACQUIRE) < 3) {  }// 等待0号线程开辟好内存空间
        const idx_t dst_off = S_up_ptr[tbeg], num_write = tl_S_upstrm.size();
        for (idx_t j = 0; j < num_write; j++)
            S_upstrm[dst_off + j] = tl_S_upstrm[j];
    }// implicit sync
    record.emplace_back(MPI_Wtime());

    idx_t num_coar = 0;
    idx_t * assum_R_rp = new idx_t [num_fine + 1];// 限制矩阵的行偏移，按最大可能空间开辟
    idx_t * assum_R_ci = new idx_t [num_fine];// 限制矩阵的列索引，它的非零元数一定是细网格的数目
    assum_R_rp[0] = 0;
    minheap.assemble();// 查询/修改前必须先完成组装（堆的建立）

    record.emplace_back(MPI_Wtime());
    // 选粗点对
    idx_t flag_heap = -1;// 现在正在被哪个线程使用

    #pragma omp parallel
    {
        int tid = omp_get_thread_num();
        // int nt  = omp_get_num_threads();
        while (true) {
            idx_t expected = -1;// 等待获得对最小堆的使用权
            while ( __atomic_compare_exchange_n( & flag_heap, & expected, tid, true, __ATOMIC_ACQUIRE, __ATOMIC_RELAXED) == false ) {
                assert(expected >= 0);
                expected = -1;// 重新修改为等待
            }
                    // __atomic_load_n( & minheap.numCounts, __ATOMIC_ACQUIRE ) > 0)
            assert(flag_heap == tid);// 现在应该是本线程在使用堆
            if (minheap.numCounts == 0) {// 堆已空，可以结束
                bool wb = __atomic_compare_exchange_n( & flag_heap, & tid, -1, true, __ATOMIC_RELEASE, __ATOMIC_RELAXED);
                assert(wb == true);
                break;
            }
            // ==================================== 以下为具体执行 ===================================
            const idx_t id = minheap.heapTop().id;// printf("id %d cnt %d\n", id, minheap.heapTop().cnt);
            assert(loc_map_f2c[id] == -1);// 之前一定还未聚合过该点

            const idx_t pbeg = S_up_ptr[id], pend = S_up_ptr[id + 1];
            idx_t aggreg_id = -1;// 准备和哪个点进行聚合？
            idx_t offset = assum_R_rp[num_coar];
            loc_map_f2c[id] = num_coar;// 更新插值表
            assum_R_ci[offset ++] = id;// 更新限制矩阵
            minheap.delMin();
            minheap.decre(pend - pbeg, & S_upstrm[pbeg]);

            for (idx_t p = pbeg; p < pend; p++) {
                idx_t tenta_id = S_upstrm[p];// tentative aggregate
                if (loc_map_f2c[tenta_id] >= 0) {// 强影响本点的这个点之前已经被聚合过了
                    // printf("Id %d try %d-th: %d but failed! has been merged to loc %d\n", id + fine_mat.beg_row, p - pbeg,
                    //     tenta_id + fine_mat.beg_row, loc_map_f2c[tenta_id]);
                    continue;// 检查下一个
                }
                else {// 找到了尚未被聚合，且确实强影响本id点的
                    aggreg_id = tenta_id;// 记录一下要聚合的点序号
                    loc_map_f2c[aggreg_id] = num_coar;// 更新插值表
                    assum_R_ci[offset ++] = aggreg_id;// 更新限制矩阵
                    minheap.delName(aggreg_id);
                    minheap.decre(S_up_ptr[aggreg_id + 1] - S_up_ptr[aggreg_id], & S_upstrm[S_up_ptr[aggreg_id]]);
                    break;// 然后跳出 本点 的循环
                }
                // if (aggreg_id == -1) printf(" failed \n");
                // else printf(" succ \n");
                // break;// 跳出 本点 的循环
            }
            // printf("Proc %d glb_F %d %d => loc_C %d\n", my_pid, id + fine_mat.beg_row, (aggreg_id == -1) ? -1 : (aggreg_id + fine_mat.beg_row), num_coar );
            num_coar ++;
            assum_R_rp[num_coar] = offset;
            // ==================================== 具体执行结束 ===================================
            // 完成该点的处理后恢复堆
            bool wb = __atomic_compare_exchange_n( & flag_heap, & tid, -1, true, __ATOMIC_RELEASE, __ATOMIC_RELAXED);
            assert(wb == true);
        }// while loop
    }// omp
    /*
    while (minheap.numCounts > 0) {
        const idx_t id = minheap.heapTop().id;// printf("id %d cnt %d\n", id, minheap.heapTop().cnt);
        assert(loc_map_f2c[id] == -1);// 之前一定还未聚合过该点

        const idx_t pbeg = S_up_ptr[id], pend = S_up_ptr[id + 1];
        idx_t aggreg_id = -1;// 准备和哪个点进行聚合？
        idx_t offset = assum_R_rp[num_coar];
        loc_map_f2c[id] = num_coar;// 更新插值表
        assum_R_ci[offset ++] = id;// 更新限制矩阵
        minheap.delMin();
        minheap.decre(pend - pbeg, & S_upstrm[pbeg]);

        for (idx_t p = pbeg; p < pend; p++) {
            idx_t tenta_id = S_upstrm[p];// tentative aggregate
            if (loc_map_f2c[tenta_id] >= 0) {// 强影响本点的这个点之前已经被聚合过了
                // printf("Id %d try %d-th: %d but failed! has been merged to loc %d\n", id + fine_mat.beg_row, p - pbeg,
                //     tenta_id + fine_mat.beg_row, loc_map_f2c[tenta_id]);
                continue;// 检查下一个
            }
            else {// 找到了尚未被聚合，且确实强影响本id点的
                aggreg_id = tenta_id;// 记录一下要聚合的点序号
                loc_map_f2c[aggreg_id] = num_coar;// 更新插值表
                assum_R_ci[offset ++] = aggreg_id;// 更新限制矩阵
                minheap.delName(aggreg_id);

                minheap.decre(S_up_ptr[aggreg_id + 1] - S_up_ptr[aggreg_id], & S_upstrm[S_up_ptr[aggreg_id]]);
                break;// 然后跳出 本点 的循环
            }
            // if (aggreg_id == -1) printf(" failed \n");
            // else printf(" succ \n");
            // break;// 跳出 本点 的循环
        }
        // printf("Proc %d glb_F %d %d => loc_C %d\n", my_pid, id + fine_mat.beg_row, (aggreg_id == -1) ? -1 : (aggreg_id + fine_mat.beg_row), num_coar );
        num_coar ++;
        assum_R_rp[num_coar] = offset;
        // printf("loc_map: %d %d %d %d %d %d %d\n", loc_map_f2c[0], loc_map_f2c[1], loc_map_f2c[2], loc_map_f2c[3], 
        //     loc_map_f2c[4], loc_map_f2c[5], loc_map_f2c[6]);

        // TODO: 
    }
    */
    assert(assum_R_rp[num_coar] == num_fine);
    assert(num_coar > 0);// 暂时不处理跨进程粗化
    delete [] S_up_ptr; S_up_ptr = nullptr;
    delete [] S_upstrm; S_upstrm = nullptr;

    record.emplace_back(MPI_Wtime());
// #define REORDER_PR
#ifdef REORDER_PR
    // 将细点编号小的往前挪，使其粗点编号也较小
    struct F2OldC {
        idx_t minF = -1;
        idx_t oldC = -1;
        bool operator < (const F2OldC & b) const {
            assert(minF != b.minF);
            return minF  < b.minF;
        }
    };
    std::vector<F2OldC> aux_arr(num_coar);
    for (idx_t oci = 0; oci < num_coar; oci++) {
        aux_arr[oci].oldC = oci;
        idx_t minF = INT_MAX;
        for (idx_t p = assum_R_rp[oci]; p < assum_R_rp[oci + 1]; p++)
            minF = MIN(minF, assum_R_ci[p]);
        aux_arr[oci].minF = minF;
    }
    std::sort(aux_arr.begin(), aux_arr.end());
    idx_t * rstr_rpt = new idx_t [num_coar + 1];
    idx_t * rstr_cid = new idx_t [num_fine];
    rstr_rpt[0] = 0;
    for (idx_t nci = 0; nci < num_coar; nci++) {
        idx_t oci = aux_arr[nci].oldC;
        for (idx_t p = assum_R_rp[oci]; p < assum_R_rp[oci + 1]; p++) {
            idx_t fi = assum_R_ci[p];
            loc_map_f2c[fi] = nci;// 重新构造细到粗的映射（相当于插值矩阵）
        }
        // 同时修改限制矩阵的行偏移
        rstr_rpt[nci + 1] = assum_R_rp[oci + 1] - assum_R_rp[oci];
    }
    for (idx_t nci = 0; nci < num_coar; nci++)
        rstr_rpt[nci + 1] += rstr_rpt[nci];
    // 修改限制矩阵的列索引
    for (idx_t nci = 0; nci < num_coar; nci++) {
        idx_t oci = aux_arr[nci].oldC;
        const idx_t num = assum_R_rp[oci + 1] - assum_R_rp[oci];
        idx_t * dst = rstr_cid + rstr_rpt[nci];
        const idx_t * src = assum_R_ci + assum_R_rp[oci];
        for (idx_t p = 0; p < num; p++)
            dst[p] = src[p];
    }
    delete [] assum_R_rp; assum_R_rp = nullptr;
    delete [] assum_R_ci; assum_R_ci = nullptr;
#else
    idx_t * rstr_rpt = assum_R_rp;
    idx_t * rstr_cid = assum_R_ci;
#endif
    // 计算粗网格矩阵
    MyParCSRMatrix<idx_t, data_t, calc_t> * coar_mat = new MyParCSRMatrix<idx_t, data_t, calc_t>(fine_mat.comm);
    idx_t * & rows_partition = coar_mat->rows_partition;
    rows_partition = new idx_t [nprocs + 1];
    rows_partition[0] = 0;
    constexpr MPI_Datatype mpi_idx_type = (sizeof(idx_t) == 4) ? MPI_INT : MPI_LONG_LONG;
    MPI_Allgather(&num_coar, 1, mpi_idx_type, rows_partition + 1, 1, mpi_idx_type, fine_mat.comm);
    for (int p = 1; p <= nprocs; p++) {
        rows_partition[p] += rows_partition[p - 1];
    }
    coar_mat->beg_row   = rows_partition[my_pid    ];
    coar_mat->end_row   = rows_partition[my_pid + 1]; assert(coar_mat->end_row - coar_mat->beg_row == num_coar);
    coar_mat->glb_nrows = rows_partition[nprocs];
#ifdef DEBUG_PRINT
    MPI_Barrier(fine_mat.comm);
    for (int p = 0; p < nprocs; p++) {
        if (p == my_pid) {
            if (p == 0) printf("Rstr in COO:\n");
            for (idx_t ci = 0; ci < num_coar; ci++) {
                // for (idx_t p = R_rp[ci]; p < R_rp[ci + 1]; p++)
                //     printf(" %d %d 1.0\n", ci + coar_mat->beg_row, R_ci[p] + fine_mat.beg_row);
                const idx_t num = R_rp[ci + 1] - R_rp[ci];
                assert(num == 2 || num == 1);
                idx_t pbeg = R_rp[ci];
                printf(" glb_C %d glb_F %d %d  \n", ci + coar_mat->beg_row,
                    R_ci[pbeg] + fine_mat.beg_row, (num == 1) ? -1 : R_ci[pbeg+1] + fine_mat.beg_row);
            }
        }
        fflush(stdout);
        MPI_Barrier(fine_mat.comm);
    }
    MPI_Barrier(fine_mat.comm);
    MPI_Abort(fine_mat.comm, -20240403);
#endif
    // 对角和非对角部分的矩阵值可以分别计算
    aggregate_mat<idx_t, data_t, calc_t>(fine_diag, num_coar, num_coar, rstr_rpt, rstr_cid, loc_map_f2c, coar_mat->diag, alpha);
    
    record.emplace_back(MPI_Wtime());

    idx_t * remote_map_f2c = new idx_t [fine_offd.ncols];
    {// 非对角部分
        // 每个进程告诉细网格上的邻居，对方在细网格上的halo区所对应的（由我负责的）列号在粗网格中的全局列号
        static_assert(sizeof(idx_t) <= sizeof(calc_t));
        
        for (idx_t s = 0; s < fine_mat.num_sends; s++) {
            const idx_t offset = fine_mat.send_cnts[s],
                        num = fine_mat.send_cnts[s + 1] - offset;
            idx_t * buf = ( idx_t *)(fine_mat.send_buf     + offset);
            const idx_t * glb_rows = fine_mat.send_row_ids + offset;
            #pragma omp parallel for schedule(static)
            for (idx_t k = 0; k < num; k++) {
                const idx_t loc_fid = glb_rows[k] - fine_mat.beg_row,
                            loc_cid = loc_map_f2c[loc_fid];
                buf[k] = loc_cid + coar_mat->beg_row;
            }
            MPI_Isend(buf, num, mpi_idx_type, fine_mat.send_pids[s], 55, fine_mat.comm, fine_mat.send_reqs + s);
        }
        for (idx_t r = 0; r < fine_mat.num_recvs; r++) {
            const idx_t offset = fine_mat.recv_cnts[r],
                        num = fine_mat.recv_cnts[r + 1] - offset;
            idx_t * buf = (idx_t *)(remote_map_f2c + fine_mat.recv_cnts[r]);
            MPI_Irecv(buf, num, mpi_idx_type, fine_mat.recv_pids[r], 55, fine_mat.comm, fine_mat.recv_reqs + r);
        }
        MPI_Waitall(fine_mat.num_sends, fine_mat.send_reqs, MPI_STATUSES_IGNORE);
        MPI_Waitall(fine_mat.num_recvs, fine_mat.recv_reqs, MPI_STATUSES_IGNORE);
#ifdef DEBUG_PRINT
        MPI_Barrier(fine_mat.comm);
        for (int p = 0; p < nprocs; p++) {
            if (p == my_pid) {
                printf("Proc %d remote_map_f2c:\n", my_pid);
                for (idx_t j = 0; j < fine_offd.ncols; j++)
                    printf("  fine loc %d glb %d Coar glb %d\n", j, fine_mat.col_map_offd[j], remote_map_f2c[j]);
                fflush(stdout);
            }
            MPI_Barrier(fine_mat.comm);
        }
#endif
        // 此时接收缓冲区中的是细矩阵非对角部分原来的局部列号在粗矩阵中的全局列号
        // 然后构造粗矩阵的col_map_offd (即局部列序号到全局列序号的映射)
        // 但可能原来的2列才对应粗矩阵中的一列，所以需要消重

        const idx_t num = fine_offd.ncols;
        std::unordered_map<idx_t, idx_t> aux_map;
        idx_t * arr = new idx_t [num];
        #pragma omp parallel for schedule(static)
        for (idx_t k = 0; k < num; k++)
            arr[k] = remote_map_f2c[k];
        // 对remote_map排序
        qsort(arr, num, sizeof(idx_t), cmpfunc<idx_t>);
        const idx_t uniq_cnt = uniquify<idx_t, data_t>(num, arr);

        // 构造coar_mat的非对角部分的列映射
        assert(coar_mat->col_map_offd == nullptr);
        coar_mat->col_map_offd = new idx_t [uniq_cnt];
        for (idx_t j = 0; j < uniq_cnt; j++) {
            coar_mat->col_map_offd[j] = arr[j];
            aux_map.emplace(arr[j], j);// 不可多线程化
        }
        // 构造非对角部分粗化所需的f2c，并储存在arr中
        #pragma omp parallel for schedule(static) // STL允许多个线程同时读
        for (idx_t f_loc = 0; f_loc < num; f_loc++) {
            const idx_t c_glb = remote_map_f2c[f_loc];
            assert(aux_map.find(c_glb) != aux_map.end());
            const idx_t c_loc = aux_map[c_glb];
            arr[f_loc] = c_loc;
        }

        aggregate_mat<idx_t, data_t, calc_t>(fine_offd, num_coar, uniq_cnt, rstr_rpt, rstr_cid, arr, coar_mat->offd, alpha);
        delete [] arr; arr = nullptr;
    }
    record.emplace_back(MPI_Wtime());
    // 至此粗网格矩阵的除通信元数据外的其它数据都已全部就绪
    
    {// 建立通信元数据
        const idx_t num_sends = fine_mat.num_sends,
                    num_recvs = fine_mat.num_recvs;
        // 在不发生跨进程粗化时，聚合型的AMG应该邻居进程数不变
        {
            coar_mat->num_sends = num_sends;
            coar_mat->send_pids = new int   [num_sends];
            coar_mat->send_cnts = new idx_t [num_sends + 1];
            idx_t * write_buf = new idx_t [fine_mat.send_cnts[num_sends]];// 按最大可能的需求开辟
            // 根据细矩阵的元数据修改得到粗矩阵的
            coar_mat->send_cnts[0] = 0;
            #pragma omp parallel for schedule(static)
            for (idx_t s = 0; s < num_sends; s ++) {
                coar_mat->send_pids[s] = fine_mat.send_pids[s];
                const idx_t offset = fine_mat.send_cnts[s];
                const idx_t * f_glbs = fine_mat.send_row_ids + offset;
                idx_t * dst_buf = write_buf + offset;

                const idx_t f_num = fine_mat.send_cnts[s + 1] - offset;
                for (idx_t fk = 0; fk < f_num; fk++) {
                    const idx_t f_loc = f_glbs[fk] - fine_mat.beg_row,
                                c_loc = loc_map_f2c[f_loc];
                    dst_buf[fk] = c_loc + coar_mat->beg_row;
                }
                qsort(dst_buf, f_num, sizeof(idx_t), cmpfunc<idx_t>);
                const idx_t uniq_cnt = uniquify<idx_t, data_t>(f_num, dst_buf);
                coar_mat->send_cnts[s + 1] = uniq_cnt;
            }
            
            for (idx_t s = 0; s < num_sends; s ++)// 单线程前缀和
                coar_mat->send_cnts[s + 1] += coar_mat->send_cnts[s];
            coar_mat->send_row_ids = new idx_t [coar_mat->send_cnts[num_sends]];
            
            #pragma omp parallel for schedule(static)
            for (idx_t s = 0; s < num_sends; s ++) {
                const idx_t f_offset = fine_mat.send_cnts[s],
                            c_offset = coar_mat->send_cnts[s];
                const idx_t c_num = coar_mat->send_cnts[s + 1] - c_offset;
                for (idx_t j = 0; j < c_num; j++)
                    coar_mat->send_row_ids[c_offset + j] = write_buf[f_offset + j];
            }
            delete [] write_buf; write_buf = nullptr;
        }
        {
            coar_mat->num_recvs = num_recvs;
            coar_mat->recv_pids = new int   [num_recvs];
            coar_mat->recv_cnts = new idx_t [num_recvs + 1];
            // 根据细矩阵的元数据修改得到粗矩阵的
            coar_mat->recv_cnts[0] = 0;
            assert(fine_offd.ncols == fine_mat.recv_cnts[num_recvs]);
            #pragma omp parallel for schedule(static)
            for (idx_t r = 0; r < num_recvs; r++) {
                coar_mat->recv_pids[r] = fine_mat.recv_pids[r];
                const idx_t offset = fine_mat.recv_cnts[r];
                const idx_t * c_glbs = remote_map_f2c + offset;
                const idx_t f_num = fine_mat.recv_cnts[r + 1] - offset;
                idx_t _poss[f_num];
                for (idx_t j = 0; j < f_num; j++)
                    _poss[j] = c_glbs[j];
                qsort(_poss, f_num, sizeof(idx_t), cmpfunc<idx_t>);
                const idx_t uniq_cnt = uniquify<idx_t, data_t>(f_num, _poss);
                coar_mat->recv_cnts[r + 1] = uniq_cnt;
            }
            for (idx_t r = 0; r < num_recvs; r++)
                coar_mat->recv_cnts[r + 1] += coar_mat->recv_cnts[r];
        }
        coar_mat->send_reqs = new MPI_Request [num_sends];
        coar_mat->recv_reqs = new MPI_Request [num_recvs];
        coar_mat->send_buf = new calc_t [coar_mat->send_cnts[num_sends]];
    }
    // printf("Proc %d : %d rows %d nnz fine => %d rows %d nnz coar\n", my_pid, num_fine, fine_diag.nnz, num_coar, coar_mat->diag->nnz);
#ifdef DEBUG_PRINT
    MPI_Barrier(fine_mat.comm);
    // for (int p = 0; p < nprocs; p++) {
    //     if (p == my_pid) {
    //         if (p == 0) printf(" AC in COO nprocs %d\n", nprocs);
    //         for (idx_t ci = 0; ci < num_coar; ci++) {
    //             for (idx_t p = coar_mat->diag->row_ptr[ci]; p < coar_mat->diag->row_ptr[ci + 1]; p++)
    //                 printf(" %d %d %.5e\n", ci + coar_mat->beg_row,
    //                     coar_mat->diag->col_idx[p] + coar_mat->beg_row, coar_mat->diag->vals[p]);
    //             for (idx_t p = coar_mat->offd->row_ptr[ci]; p < coar_mat->offd->row_ptr[ci + 1]; p++)
    //                 printf(" %d %d %.5e\n", ci + coar_mat->beg_row,
    //                     coar_mat->col_map_offd[ coar_mat->offd->col_idx[p] ], coar_mat->offd->vals[p]);
    //         }
    //         fflush(stdout);
    //     }
    //     MPI_Barrier(fine_mat.comm);
    // }
    for (int p = 0; p < nprocs; p++) {
        if (p == my_pid) {
            if (p == 0) printf(" Rstr in COO nprocs %d\n", nprocs);
            for (idx_t ci = 0; ci < num_coar; ci++) {
                for (idx_t p = rstr_rpt[ci]; p < rstr_rpt[ci + 1]; p++)
                    printf(" %d", rstr_cid[p] + fine_mat.beg_row);
                printf("\n");
            }
            fflush(stdout);
        }
        MPI_Barrier(fine_mat.comm);
    }
    MPI_Barrier(fine_mat.comm);
    MPI_Abort(fine_mat.comm, -20240402);

    // const idx_t c_nrows = rows_partition[nprocs],
    //             c_beg = rows_partition[my_pid],
    //             c_end = rows_partition[my_pid + 1];
    // MyParCSRVector<idx_t, calc_t> test_x(coar_mat->comm, c_nrows, c_beg, c_end, nullptr);
    // MyParCSRVector<idx_t, calc_t> test_y(test_x);
    // srand(1993);
    // for (idx_t i = 0; i < test_x.end_row - test_x.beg_row; i++) {
    //     test_x.data[i] = (calc_t) rand() / (calc_t) RAND_MAX;
    // }
    // MPI_Barrier(fine_mat.comm);
    // for (int p = 0; p < nprocs; p++) {
    //     if (p == my_pid) {
    //         if (p == 0) printf("\ntest_x :\n");
    //         for (idx_t ci = 0; ci < num_coar; ci++)
    //             printf(" %.15e\n", test_x.data[ci]);
    //     }
    //     fflush(stdout);
    //     MPI_Barrier(fine_mat.comm);
    // }
    // coar_mat->Mult(1.0, test_x, 0.0, test_y);
    // double ck_dot = vec_dot<idx_t, calc_t>(test_y, test_y);
    // if (my_pid == 0) {
    //     printf("\nck_dot = %.15e\n", ck_dot);
    //     fflush(stdout);
    // }
#endif
    delete [] remote_map_f2c; remote_map_f2c = nullptr;
    record.emplace_back(MPI_Wtime());
    if (my_pid == 0) {
        for (size_t i = 0; i < record.size() - 1; i++)
            printf(" %.5f ", record[i+1] - record[i]);
        printf("\n");
    }

    // 返回的所有粗化信息
    CoarsenWrapper<idx_t, data_t, calc_t> * ret = new CoarsenWrapper<idx_t, data_t, calc_t>;
    ret->loc_map_f2c = loc_map_f2c;
    ret->R_rp = rstr_rpt;
    ret->R_ci = rstr_cid;
    ret->coar_mat = coar_mat;

    return ret;
}

#endif